# -*- encoding: utf-8 -*-
'''
@File    :   patch.py
@Time    :   2021/11/30 15:25
@Author  :   ZhangChaoYang
@Desc    :   把图像划分成很多小Patch，再从patch还原回图像
'''

import tensorflow as tf

from tensorflow.keras import Model
from tensorflow.keras.layers import Reshape


class Patches(Model):
    def __init__(self, patch_size):
        super(Patches, self).__init__()
        self.patch_height, self.patch_width = patch_size
        self.resize = Reshape(target_shape=(-1, self.patch_height * self.patch_width))

    def call(self, inputs, training=None, mask=None):
        '''把一张图片分割成很多patch
        @param inputs: 多张图片
        '''
        patches = tf.image.extract_patches(
            images=inputs,
            sizes=[1, self.patch_height, self.patch_width, 1],
            strides=[1, self.patch_height, self.patch_width, 1],
            rates=[1, 1, 1, 1],
            padding='VALID',
        )  # patches是4维(batch_size,几行patch,几列patch,一个patch内有几个像素)
        patches = self.resize(patches)  # patches是3维(batch_size,一共有几个patch,一个patch内有几个像素)
        return patches

    def reconstruct_from_patch(self, patches, img_width):
        '''把很多patch还原为一张图片
        @param patches: 一张图片对应的多个patch
        '''
        num_patches = patches.shape[1]
        patches = tf.reshape(patches, shape=(-1, num_patches, self.patch_height, self.patch_width))
        n = img_width // self.patch_width
        rows = tf.split(patches, n, axis=1)
        rows = [tf.concat(tf.unstack(x, axis=1), axis=1) for x in
                rows]  # unstack把axis=1这一维提出来一个Tensor变成一个list，再通过concat把list变成一个Tensor。所以先过unstack再过concat实际上相当于是合并了axis=1和axis=2这两维
        reconstructed = tf.concat(rows, axis=2)
        return reconstructed


if __name__ == '__main__':
    img_height, img_width = 20, 10
    images = tf.ones(shape=(1, img_height, img_width, 1))
    patch_height, patch_width = 6, 5
    model = Patches(patch_size=(patch_height, patch_width))
    patches = model(images)
    print(patches.shape)  # (batch_size,一共有几个patch,一个patch内有几个像素)。(1, 6, 30)，20不能被6整除，最后2行会被丢弃掉
    reconstructed_img = model.reconstruct_from_patch(patches, img_width)
    print(reconstructed_img.shape)

# python .\models\patch.py
